#include "manager.h"
#include "simulatorParameters.h"

struct edgeInfo
{
	int I;
	int J;
	double *a;
};

void initializeMatrices(double **edgeCenters, double **Dij, double **Cij,
						int **edges_IJ);
void assemblyMatrixEdges(int e, pEdge edge,  double **Dij,  double **Cij,
						 int **edges_IJ, const double *K, const double &A);
bool isEdgeOnPartitionBdry(pMesh theMesh, pEdge e);
int unifyEdgeCoefficientsMatrices(pMesh theMesh);

#ifdef _HAVE_PETSC_

int getUnifiedCoefficients(pMesh theMesh, Mat *global_A, list<edgeInfo> &bdryEInfoList);

#endif

void updateEdgeCoefficients(pMesh theMesh, list<edgeInfo> &bdryEInfoList);
void exportCoefficients(pMesh theMesh);
void printData(ofstream &fid, pEntity entity, const char **str, const int &n);

// a loop over all elements is made and for each one the edge matrix is computed
// edge matrix is 2x2
// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
// IN:
// theMesh - mesh object
// OUT:
void calculateCoefficients2D_new(pMesh theMesh, set<int> &setOfDomains, const char *meshFilename)
{
	if ( theMesh->getDim() 	!= 2 )
		throw Exception(__LINE__,__FILE__,"Three dimensional meshes are not supported\n");

	cout << "calculateCoefficients2D_new\n";
	unsigned short int i,j,numEdges;
	int *edges_IJ[3];
	double *edgeCenters[3], *Dij[3], *Cij[3], I[3], J[3];

	char *Cij_string[6] = {"Cx0","Cy0","Cx1","Cy1","Cx2","Cy2"};
	char *Dij_string[6] = {"Dx0","Dy0","Dx1","Dy1","Dx2","Dy2"};

	SimulatorParameters fm;

	fm.loadParameters();

	// allocate auxiliaries matrices
	for (i=0; i<3; i++)
	{
		edgeCenters[i] = new double[3];
		edges_IJ[i] = new int[2];
		Cij[i] = new double[2];
		Dij[i] = new double[2];
	}
	pGEntity gEntity;
	vector<pVertex> vertex;
	pEdge edge[3];
	double sum[3] = {.0, .0, .0};

	// loop over elements:
	// for each one, calculate all data needed and attach them to the its edges
	// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =

	pEntity face;
	FIter fit = M_faceIter(theMesh);
	while (face = FIter_next(fit))
	{
		gEntity =(pGEntity)F_whatIn((pFace)face);				// get element flag
		double *K = 0;//fm.getKTensor( GEN_tag(gEntity) );		// interpret element flag
		double faceCenter[3] = {.0, .0, .0};
		F_center(face,faceCenter);								// element centroid

		initializeMatrices(edgeCenters,Dij,Cij,edges_IJ);	// zeros entries

		double A = F_area(face)/3.0;							// element area
		M_GetVertices(face,vertex);							// extract face's vertices
		// assign to all face's vertices part of control volume
		for (i=0; i<3; i++)
		{
			double nvol = .0;
			EN_getDataDbl(vertex[i],MD_lookupMeshDataId("NodeVolume"),&nvol);
			nvol += A;
			EN_attachDataDbl(vertex[i],MD_lookupMeshDataId("NodeVolume"),nvol);
		}
		vertex.clear();
		// loop over all three face's edges
		for (i=0; i<3; i++)
		{
			edge[i] = F_edge(face, i);

			E_center(edge[i],edgeCenters[i]);				// edge centroid

			M_GetVertices(edge[i],vertex);
			V_coord(vertex[0],I);
			V_coord(vertex[1],J);

			// edge vector, used as a reference vector
			double IJ[2] = {J[0]-I[0], J[1]-I[1]};

			// vector IJ must point from the smaller vertex ID to the greater
			edges_IJ[i][0] = EN_id(vertex[0]);
			edges_IJ[i][1] = EN_id(vertex[1]);
			if ( edges_IJ[i][0] > edges_IJ[i][1] )
			{
				swap(edges_IJ[i]);
				for (j=0; j<2; j++) IJ[j] = -IJ[j];
			}

			// Dij is orthogonal to IJ
			Dij[i][0] = IJ[1]/2.0;
			Dij[i][1] = -IJ[0]/2.0;

			// vector: from element center to edge middle point, used as a reference vector
			double v[2] = {edgeCenters[i][0]-faceCenter[0],edgeCenters[i][1]-faceCenter[1]};

			// Dij must point to outside element
			if (  dot(Dij[i],v) <= 0.0 )
				for (j=0; j<2; j++) Dij[i][j] = -Dij[i][j];

			// Cij is orthogonal to v
			Cij[i][0] = v[1];
			Cij[i][1] = -v[0];

			// Cij must point as if I inside the CV and J outside
			if ( dot(Cij[i],IJ) <= 0.0 )
			{
				for (j=0; j<2; j++)
					Cij[i][j] = -Cij[i][j];
			}
			vertex.clear();

		}

		// calculated all Dij and Cij for all edges of a specific face,
		// edge's matrices must be assembled
		for (i=0; i<3; i++)
			assemblyMatrixEdges(i, edge[i], Dij, Cij, edges_IJ, K, A);

		// associate Dij from edges to element
		int k = 0;
		for (i=0; i<3; i++)
			for (j=0; j<2; j++)
			{
				EN_attachDataDbl(face,MD_lookupMeshDataId(Cij_string[k]),Cij[i][j]);
				EN_attachDataDbl(face,MD_lookupMeshDataId(Dij_string[k]),Dij[i][j]);
				k++;
			}
	}
	FIter_delete(fit);
	// deallocate memory
	// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
	for (i=0; i<3; i++) delete edgeCenters[i];

#ifdef _HAVE_PETSC_
	// for parallel simulation. Does nothing in serial
	unifyEdgeCoefficientsMatrices(theMesh);
#endif

	exportCoefficients(theMesh);
}

// for each element, all three edges matrices are stored.
// A matrix 2x2 for each edge
// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
// IN:
// e - one of the three face's edges
// edge -  one of the three face's edges
// Dij - matrix with Dij coefficients for all three edges
// Cij - matrix with Cij coefficients for all three edges
// edges_IJ - matrix with IJ node indices for all three edges
// K - element physical property
// A - face area
// OUT:
void assemblyMatrixEdges(int e, pEdge edge,  double **Dij,  double **Cij,
						 int **edges_IJ, const double *K, const double &A)
{
	double p = 1.0/6.0;
	double a[4] = { .0, .0, .0, .0};
	char *str[4] = {"a_00", "a_01", "a_10", "a_11"};
	int i,j;

	for (i=0; i<4; i++)
		EN_getDataDbl(edge,MD_lookupMeshDataId(str[i]),&a[i]);

	// search for edge opposite to node I
	double *dij;
	int I = edges_IJ[e][0];
	for (i=0; i<3; i++)
	{
		if ( edges_IJ[i][0] != I && edges_IJ[i][1] != I )
		{
			dij = Dij[i];
			break;
		}
	}

	// calculate a[0][0]
	double *cij = Cij[e];
	double W_kj[2];
	multTensorVector2D(K,dij,W_kj);
	W_kj[0] /= A;
	W_kj[1] /= A;

	a[0] += p*2.0*dot(W_kj,cij);

	// calculate a[0][1]
	double aux[2] = { .0, .0};
	for (i=0; i<2; i++)	aux[i] = 3.0*Cij[e][i] + Dij[e][i];
	a[1] += -p*dot(W_kj,aux);

	// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
	// search for edge opposite to node J
	int J = edges_IJ[e][1];
	for (j=0; j<3; j++)
	{
		if ( edges_IJ[j][0] != J && edges_IJ[j][1] != J )
		{
			dij = Dij[j];
			break;
		}
	}
	double W_ki[2];
	multTensorVector2D(K,dij,W_ki);
	W_ki[0] /= A;
	W_ki[1] /= A;

	// calculate a[1][0]. NOTE: Matrix a is symmetric => a[0][1] = a[1][0]
	a[2] = a[1];

	// calculate a[1][1]
	a[3] += -p*2.0*dot(W_ki,Cij[e]);

	// assign all four values of matriox a to edge
	for (i=0; i<4; i++)
		EN_attachDataDbl(edge,MD_lookupMeshDataId(str[i]),a[i]);

}

// zero matrices entries
// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
// IN:
// edgeCenters - matrix with edges centroids
// Dij - matrix with Dij coefficients for all three edges
// Cij - matrix with Cij coefficients for all three edges
// edges_IJ - matrix with IJ node indices for all three
// OUT:
void initializeMatrices(double **edgeCenters, double **Dij, double **Cij, int **edges_IJ)
{
	for (int i=0; i<3; i++)
		for (int j=0; j<2; j++)
		{
			Dij[i][j] = 0.0;
			Cij[i][j] = 0.0;
			edgeCenters[i][j] = 0.0;
			edges_IJ[i][j] = 0;
		}
}

#ifdef _HAVE_PETSC_
// For parallel simulation only!
// An edge matrix contains data from all elements (for 2D triangles)
// surrounding it. For those edges on partion boundary, the coeffient matrix
// is incomplete. After calculation had been made, an edge on partition boundary
// must has only one coefficient matrix. So, processor A sends its values (for a
// specific edge) to processor B and B add the received value to the current (for
// the same edge). Processor B does the same as A and A adds the value.
// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
// IN:
// theMesh - mesh object
// OUT:
int unifyEdgeCoefficientsMatrices(pMesh theMesh)
{
	if ( P_size() == 1) return 1;

	int i;
	int n = M_getMaxVId(theMesh);
	int d_nz = 20, o_nz = 20;
	PetscErrorCode ierr;

	// a global matrix A is used to unify remote edges matrices
	// edge's matrix is 2x2 => 4 scalars, so 4 matrices A are used
	Mat A[4];
	for (i=0; i<4; i++)
		ierr = MatCreateMPIAIJ(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,
    							n+1,n+1,
    							d_nz,PETSC_NULL,
    							o_nz,PETSC_NULL,
    							&A[i]); CHKERRQ(ierr);

    vector<pVertex> vertex;

	// list to store all IDs from edges with remote copies
	list<edgeInfo> bdryEInfoList;
	list<edgeInfo>::iterator iter;
	// get matrix coefficients
	double *a;

	// loop over all edges and take only those with remote copies
	EIter eit = M_edgeIter(theMesh);
	while ( pEntity e = EIter_next(eit) )
	{
		if ( isEdgeOnPartitionBdry(theMesh,e) )
		{
			a = new double[4];
			char *str[4] = {"a_00", "a_01", "a_10", "a_11"};
			for (i=0; i<4; i++)
			{
				a[i] = 0.0;
				EN_getDataDbl(e,MD_lookupMeshDataId(str[i]),&a[i]);
			}

			// get edge vertex IDs
			M_GetVertices(e,vertex);
			int IJ[2];
			IJ[0] = EN_id(vertex[0]);
			IJ[1] = EN_id(vertex[1]);
			if ( IJ[0]>IJ[1] ) swap(IJ);

			edgeInfo einfo;
			einfo.I = IJ[0];
			einfo.J = IJ[1];
			einfo.a = a;
			bdryEInfoList.push_back(einfo);

			for (i=0; i<4; i++)
				ierr = MatSetValue(A[i],IJ[0],IJ[1],a[i],ADD_VALUES);

			vertex.clear();
		}
	}
	EIter_delete(eit);

	// assembly all matrices
	for (i=0; i<4; i++)
	{
		ierr = MatAssemblyBegin(A[i],MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
		ierr = MatAssemblyEnd(A[i],MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
	}

	// after matrices As has been assembled, it's time to gather the unified
	// edge's matrices an store them on a list
	getUnifiedCoefficients(theMesh, A, bdryEInfoList);

	updateEdgeCoefficients(theMesh,bdryEInfoList);

	// clean up all unnecessary stuffes
	for (i=0; i<4; i++){
		ierr = MatDestroy(A[i]); CHKERRQ(ierr);
	}
	for (iter=bdryEInfoList.begin(); iter != bdryEInfoList.end(); iter++)
	{
		edgeInfo einfo = (*iter);
		delete[] einfo.a;
	}
	bdryEInfoList.clear();
}

// get all local edges with remote copies matrices values
// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
// IN:
// theMesh - mesh object
// global_A - Petsc matrices with edge's matrix
// bdryEInfoList - list to stores edge's matrices
// OUT:
int getUnifiedCoefficients(pMesh theMesh, Mat *global_A, list<edgeInfo> &bdryEInfoList)
{
	int i;
	PetscErrorCode ierr;

	// global_A é usada para unificar os coeficientes das arestas. Apos a unificaçao,
	// é necessario usar uma outra matrix, local_A, para que cada processador possa
	// recuperara os valores das suas respectivas arestas. Um processador não pode
	// obter valores de global_A que não são locais. è preciso usar MatGetSubMatrixRaw
	// informando a esta função que linhas de global_A deseja.
	// se um processador necessita de valores que estão nas linhas 34,55,66, por exemplo,
	// apos chamar MatGetSubMatrixRaw, estas linhas estarão acessiveis localmente nas linhas 0,1,2
	// Logo, é preciso mapear as linhas para que se possa pegar o valor correto.

	// store all local I edge node index (Is nodes refers to matrix rows) with
	// remote copies
	set<int> setRows;
	list<edgeInfo>::iterator listIter;
	for (i=0, listIter=bdryEInfoList.begin(); listIter != bdryEInfoList.end(); listIter++)
		setRows.insert(listIter->I);

	// map all I node indices (maybe the previous loop could be eliminated)
	set<int>::iterator setIter;
	map<int,int> mapRows;
	for (i=0, setIter=setRows.begin(); setIter != setRows.end(); setIter++)
		mapRows.insert( pair<int,int>(*setIter,i++));

	setRows.clear();

	int nrows = mapRows.size();
	int *rows = new int[nrows];
	map<int,int>::iterator mapIter;
	for (i=0, mapIter=mapRows.begin(); mapIter != mapRows.end(); mapIter++)
		rows[i++] = mapIter->first;

	int ncols = M_getMaxVId(theMesh)+1;
	int *cols = new int[ncols];
	for (i=0; i<ncols; i++) cols[i] = i;

	Mat local_A[4];
	// processor p must get values that are owned to other processors
	for (int i=0; i<4; i++){
		ierr = MatGetSubMatrixRaw(global_A[i], nrows, rows,
							ncols, cols, PETSC_DECIDE,
							MAT_INITIAL_MATRIX, &local_A[i]);  CHKERRQ(ierr);
	}
	int mowner, nowner;
	ierr = MatGetOwnershipRange(local_A[0],&mowner,&nowner);

	// transfer required matrices values from global matrix to a list
	for (listIter=bdryEInfoList.begin(); listIter != bdryEInfoList.end(); listIter++)
	{
		mapIter = mapRows.find(listIter->I);
		int col = listIter->J;
		int row = mapIter->second + mowner;

		double val = 0.0;
		for (i=0; i<4; i++){
			ierr = MatGetValues(local_A[i],	1,&row,	1,&col,	&val);CHKERRQ(ierr);
			listIter->a[i] = val;
		}
	}
	// clean up all unnecessary stuffs
	for (i=0; i<4; i++){
		ierr = MatDestroy(local_A[i]); CHKERRQ(ierr);
	}
	delete[] rows;
	delete[] cols;
}
#endif //-D_HAVE_PETSC_

// check if edge has remote copies
// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
// IN:
// theMesh - mesh object
// e - edge object
// OUT:
// true or false
bool isEdgeOnPartitionBdry(pMesh theMesh, pEdge e)
{
	return ( M_numRemoteCopies(theMesh,e) > 0 );
}

// update edge's matrices
// = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
// IN:
// theMesh - mesh object
// bdryEInfoList - list to stores edge's matrices
// OUT:

void updateEdgeCoefficients(pMesh theMesh, list<edgeInfo> &bdryEInfoList)
{
	char *str[4] = {"a_00", "a_01", "a_10", "a_11"};

	// loop over all edges and take only those with remote copies
	EIter eit = M_edgeIter(theMesh);
	list<edgeInfo>::iterator listIter = bdryEInfoList.begin();
	while ( pEntity e = EIter_next(eit) )
	{
		if ( isEdgeOnPartitionBdry(theMesh,e) )
		{
			for (int i=0; i<4; i++)
				EN_attachDataDbl(e,MD_lookupMeshDataId(str[i]),listIter->a[i]);
			listIter++;
		}
	}
	EIter_delete(eit);
}

void exportCoefficients(pMesh theMesh)
{
	cout << "exporting edge coefficients\n";

	ofstream fid;
	fid.open("geo-coeff-FVM2D-new.dat");

	fid << setprecision(16) << fixed << scientific;

	const char *str1[4] = {"a_00", "a_01", "a_10", "a_11"};
	pEntity vertex, edge, face;

	// print edge matrix
	fid << M_numEdges(theMesh) << endl;
	EIter eit = M_edgeIter(theMesh);
	while ( edge = EIter_next(eit) )
		printData(fid, edge, str1, 4);
	EIter_delete(eit);

	// print nodal volumes
	const char *volume_str[1] = {"NodeVolume"};
	VIter vit = M_vertexIter(theMesh);
	while (vertex = VIter_next(vit))
		printData(fid,vertex,volume_str	,1);
	VIter_delete(vit);

	// print Dij's associated to elements' edge
	const char *Cij_string[6] = {"Cx0","Cy0","Cx1","Cy1","Cx2","Cy2"};
	const char *Dij_string[6] = {"Dx0","Dy0","Dx1","Dy1","Dx2","Dy2"};
	FIter fit = M_faceIter(theMesh);
	while (face = FIter_next(fit))
	{
		printData(fid,face,Cij_string,6);
		printData(fid,face,Dij_string,6);
	}
	FIter_delete(fit);

	fid.close();
}

void printData(ofstream &fid, pEntity entity, const char **str, const int &n)
{
	for (int i=0; i<n; i++)
	{
		double val = .0;
		EN_getDataDbl(entity,MD_lookupMeshDataId(str[i]),&val);
		fid << val << " ";
		//cout << val << " ";
	}
	fid << endl;
	//cout<<endl;
}
